import tkinter as tk
from tkinter import filedialog, messagebox, ttk
import os
import time
from concurrent.futures import ThreadPoolExecutor
import requests
import m3u8
from urllib.parse import urljoin
import threading
from datetime import datetime

""" todo:
1. 点击listview，修改单元格的值 2025-01-16
2. listview的增删改 2025-01-16
3. ui界面增加console
4. .m3u8直播接口的导入
5. 重复/类似链接的测试
6. 最低下载速度的设置
7. 非实时链接的测试
8. 排序的优化，比如字符串排序时，拆分成数值部分和字符串部分，先按字符串部分排序，再按数值部分排序。

"""
# 让msgbox显示在当前窗口最中间


class LiveTestApp:
    def __init__(self, root):
        self.root = root
        self.root.title("直播源测试工具")
        self.input_file_path = tk.StringVar()
        self.create_widgets()
        self.create_menu()
        # 获取当前脚本的目录
        self.db = self.read_file_to_dict("db.table")
        self.show_db_data()
        # 绑定双击事件
        self.tree.bind("<Double-1>", self.on_double_click)
        # 绑定窗口关闭事件
        self.threads = []
        self.root.protocol("WM_DELETE_WINDOW", self.on_closing)
        # 初始化排序状态
        self.sort_column = None
        self.sort_reverse = False
        # 绑定表头点击事件
        for col in self.tree["columns"]:
            self.tree.heading(col, command=lambda c=col: self.sort_treeview(c))
        self.message_history = []  # 用于存储消息历史的列表
        self.console_print("欢迎使用直播源测试工具！作者：czwziy@www.ynlo.top")

    def on_closing(self):
        self.running = False
        # 保存db_dict到文件
        self.save_dict_to_db()
        self.root.destroy()

    def create_widgets(self):
        # 按钮组frame
        self.button_frame = tk.Frame(
            self.root,
            highlightthickness=1,
            highlightcolor="blue",
            highlightbackground="black",
        )
        self.button_frame.grid(
            row=1, column=0, columnspan=5, padx=10, pady=10, sticky="nsew"
        )
        tk.Button(
            self.button_frame, width=10, text="切换主数据", command=self.show_db_data
        ).grid(row=0, column=0, padx=10, pady=20)
        # 启动按钮
        tk.Button(self.button_frame, width=10, text="清除无效链接").grid(
            row=0, column=1, padx=10, pady=20
        )   
        # 测试按钮
        tk.Button(self.button_frame, width=10, text="测试",command=self.dev_test).grid(
            row=0, column=2, padx=10, pady=20
        )

        #dev_test
        tk.Button(
            self.root, width=10, text="开始测试", command=self.start_test
        ).grid(row=1, column=4,  padx=10, pady=20, sticky="ens")
        # Treeview 和 Scrollbar
        self.tree = ttk.Treeview(
            self.root,
            columns=(
                "序号",
                "直播源名称",
                "直播源地址",
                "直播源地区",
                "直播源延迟",
                "平均延迟",
                "直播源速度",
                "平均速度",
                "测试次数",
                "测试日期",
            ),
            show="headings",
            height=30,
        )
        self.tree.grid(row=3, column=0, columnspan=5, padx=10, pady=10, sticky="nsew")
 
        # 添加信息提示的Label
        self.message_label = tk.Label(
            self.root,
            text="",
            height=15,
            borderwidth=1,  # 添加边框
            relief="solid",  # 设置边框样式
            bg="black",  # 设置背景颜色
            fg="white",  # 设置前景颜色
            anchor="nw",  # 设置文本对齐方式
        )
        self.message_label.grid(row=4, column=0, columnspan=5, padx=10, pady=10, sticky="nsew")

        # 设置列标题
        for col in self.tree["columns"]:
            self.tree.heading(col, text=col)

        # 设置列宽
        self.tree.column("序号", width=50)
        self.tree.column("直播源名称", width=100)
        self.tree.column("直播源地址", width=300)
        self.tree.column("直播源地区", width=100)
        self.tree.column("直播源延迟", width=100)
        self.tree.column("直播源速度", width=100)
        self.tree.column("测试日期", width=150)

        scrollbar = ttk.Scrollbar(self.root, orient="vertical", command=self.tree.yview)
        scrollbar.grid(row=3, column=5, sticky="ns")
        self.tree.config(yscrollcommand=scrollbar.set)

        # 设置宽度
        self.root.columnconfigure(0, minsize=10)
        self.root.resizable(False, False)
 
    # 双击treeview，修改单元格值
    def on_double_click(self, event):
        # 获取点击的行和列
        region = self.tree.identify_region(event.x, event.y)
        if region == "cell":
            item = self.tree.identify_row(event.y)
            column = self.tree.identify_column(event.x)
            column_index = int(column[1:]) - 1  # 列索引从1开始，需要减1
            # 获取当前单元格的值
            current_value = self.tree.item(item, "values")[column_index]
            # 创建一个Entry控件进行编辑
            x, y, width, height = self.tree.bbox(item, column)
            entry = tk.Entry(self.tree, width=width // 8)
            entry.place(x=x, y=y, height=height, width=width)

            # 设置Entry的初始值
            entry.insert(0, current_value)
            entry.focus_set()

            def on_entry_return(event):
                new_value = entry.get()
                # 更新Treeview中的值
                current_values = list(self.tree.item(item, "values"))
                current_values[column_index] = new_value
                self.tree.item(item, values=current_values)

                # 更新db_dict中的值
                link = current_values[2]  # 假设链接在第2列
                self.db[link][
                    column_index - 1
                ] = new_value  # column_index + 1 因为db_dict中没有序号列

                # 销毁Entry控件
                entry.destroy()

            entry.bind("<Return>", on_entry_return)
            entry.bind("<FocusOut>", on_entry_return)

    def create_menu(self):
        # 创建菜单栏
        menubar = tk.Menu(self.root)
        self.root.config(menu=menubar)

        # 添加文件菜单
        file_menu = tk.Menu(menubar, tearoff=0)
        file_menu.add_command(label="导入文件到库", command=self.browse_file_to_db_file)
        file_menu.add_command(label="导出接口文件", command=self.export_interface_file)
        file_menu.add_separator()
        file_menu.add_command(label="退出", command=self.root.quit)
        menubar.add_cascade(label="文件", menu=file_menu)

        # 添加编辑菜单
        edit_menu = tk.Menu(menubar, tearoff=0)
        edit_menu.add_command(label="切换数据", command=self.show_db_data)
        edit_menu.add_command(label="清除无效链接")
        menubar.add_cascade(label="编辑", menu=edit_menu)

        # 添加帮助菜单
        help_menu = tk.Menu(menubar, tearoff=0)
        help_menu.add_command(
            label="关于",
            command=lambda: messagebox.showinfo(
                "关于",
                "直播源测试工具 v1.0\n作者：czwziy@www.ynlo.top",
                parent=self.root,
            ),
        )
        menubar.add_cascade(label="帮助", menu=help_menu)
    # 打印
    def console_print(self, msg):
        # 将新消息添加到消息历史列表中
        self.message_history.append(msg)

        # 如果消息历史列表超过14行，则删除最早的消息
        if len(self.message_history) > 14:
            self.message_history.pop(0)

        # 更新message_label的文本内容
        self.message_label.config(text="\n".join(self.message_history))
        self.root.update()


    # 读取文件，返回链接唯一的字典
    def read_file_to_dict(self, file_path):
        result = {}
        if os.path.exists(file_path):
            with open(file_path, "r", encoding="utf-8") as file:
                lines = file.readlines()
                for line in lines:
                    # line的格式是：直播源名称\t直播源地址\t直播源地区\t直播源延迟\t平均延迟\t直播源速度\t平均延迟\t测试次数\t最近测试日期
                    parts = line.rstrip("\n").split("\t")
                    if len(parts) >= 3:
                        result[parts[1]] = parts
        return result

    # 转换链接
    def convert_links(self, lines):
        result = {}
        region = None
        for line in lines:
            parts = line.strip().split(",")
            if len(parts) > 1:
                channel_name = parts[0].strip()
                url = parts[1].strip()
                if "#genre#" in url:
                    region = channel_name
                    continue
                urls = url.strip().split("#")
                for link in urls:
                    if not result.get(link.strip()):
                        result[link] = {
                            "channel_name": channel_name,
                            "region": region,
                        }
        return result

    # 导入文件到数据库
    def browse_file_to_db_file(self):
        file_path = filedialog.askopenfilename(filetypes=[("Text files", "*.*")])
        if file_path:
            with open(file_path, "r", encoding="utf-8") as file:
                lines = file.readlines()
                # 默认单仓 #genre#
                i = 0
                new_links = {}
                links = self.convert_links(lines)
                for link, info in links.items():
                    if not link.startswith("http"):
                        continue
                    if not self.db.get(link):
                        i += 1
                        arr = [
                            info["channel_name"],
                            link,
                            info["region"],
                            "",
                            "",
                            "",
                            "",
                            "",
                            "",
                        ]
                        self.db[link] = arr
                        new_links[link] = arr
            # 追加到db文件
            self.save_dict_to_db(new_links)
            # 更新树视图
            self.show_db_data()
            messagebox.showinfo("提示", f"导入文件成功！新增{i}个源")

    # 保存数据到db.table中
    def save_dict_to_db(self, new_links=None):
        if new_links:
            with open("db.table", "a", encoding="utf-8") as file:
                for info in new_links.values():
                    item = "\t".join(info)
                    file.write(item + "\n")
        else:
            with open("db.table", "w", encoding="utf-8") as file:
                for info in self.db.values():
                    item = "\t".join(info)
                    file.write(item + "\n")

    # 展示db数据
    def show_db_data(self):
        # 清空Treeview
        for item in self.tree.get_children():
            self.tree.delete(item)
        # 读取db_dict中的数据
        for index, info in enumerate(self.db.values(), start=1):
            self.tree.insert("", tk.END, values=(index,) + tuple(info))

        # 排序Treeview

    def sort_treeview(self, column):
        # 获取当前列的索引
        column_index = self.tree["columns"].index(column)

        # 获取所有行的数据
        items = [(self.tree.set(k, column), k) for k in self.tree.get_children("")]

        # 根据数据类型进行排序
        if column in [
            "序号",
            "直播源延迟",
            "平均延迟",
            "直播源速度",
            "平均速度",
            "测试次数",
        ]:
            # 数值类型的列
            items.sort(
                key=lambda x: float(x[0] if x[0] else 0), reverse=self.sort_reverse
            )
        else:
            # 字符串类型的列
            items.sort(key=lambda x: x[0], reverse=self.sort_reverse)

        # 更新Treeview
        for index, (val, k) in enumerate(items):
            self.tree.move(k, "", index)

        # 更新排序状态
        self.sort_column = column
        self.sort_reverse = not self.sort_reverse

        # 更新表头箭头
        self.update_heading_arrow(column)

    # 更新表头箭头
    def update_heading_arrow(self, column):
        # 清除所有表头的箭头
        for col in self.tree["columns"]:
            self.tree.heading(
                col, text=col, command=lambda c=col: self.sort_treeview(c)
            )

        # 添加箭头
        if self.sort_reverse:
            arrow = "↓"
        else:
            arrow = "↑"
        self.tree.heading(
            column,
            text=f"{column} {arrow}",
            command=lambda c=column: self.sort_treeview(c),
        )

    # 获取第一个片段链接
    def get_seg_link(self, m3u8_url):
        m3u8_obj = m3u8.load(m3u8_url, timeout=1)
        if m3u8_obj:
            first_seg_url = m3u8_obj.segments[0].uri
        return urljoin(m3u8_url, first_seg_url)

    # 测试链接速度
    def test_m3u8_connectivity(self, m3u8_url):
        try:
            # 尝试加载 m3u8 文件
            start_time = time.time()
            m3u8_first_url = self.get_seg_link(m3u8_url)
            response = requests.get(m3u8_first_url, timeout=1)
            if response.status_code == 200:
                file_size = len(response.content)
                end_time = time.time()
                time_len = end_time - start_time
                speed = file_size / time_len
                return [round(time_len, 0), round(speed / 1024, 0)]  # 返回延迟和速度
            else:
                return None
        except Exception as e:
            return None

    # 开始测试
    def start_test(self):
        def test_task():
            with ThreadPoolExecutor(max_workers=10) as executor:
                future_to_line = {}
                i = 0
                for item in self.db.values():
                    if not self.running:
                        return
                    if not ".m3u8" in item[1] or ".m3u8." in item[1]:
                        self.console_print(f"非直播源：{item[1]}")
                        continue
                    if item[1].startswith("http") and item[4] == "":
                        i += 1
                        future = executor.submit(self.test_m3u8_connectivity, item[1])
                        future_to_line[future] = (item[1], future)
                self.console_print(f"开始测试:共导入{i}个链接")
                i = 0
                for future, (link, _) in future_to_line.items():
                    i += 1
                    self.console_print(f"开始测试{i}：{link}")
                    if not self.running:
                        return
                    try:
                        result = future.result()
                        if result is not None:
                            self.console_print(f"测试成功{i}：{link}")
                            with self.db_lock:  # 使用锁来保护共享资源
                                self.db[link][4] = result[0]
                                self.db[link][6] = result[1]
                        else:
                            with self.db_lock:  # 使用锁来保护共享资源
                                self.db[link][4] = -1
                                self.db[link][6] = -1
                            self.console_print(f"测试失败{i}：{link}")
                        if i % 20 == 0:
                            self.console_print("保存数据")
                            self.root.after(100, self.show_db_data)
                            self.root.after(100, self.save_dict_to_db)
                    except Exception as e:
                        self.console_print(f"Error testing link {link}: {e}")

            executor.shutdown(wait=True)  # 确保所有任务都已完成
            self.console_print("所有测试完成")

        self.console_print("开始测试")
        self.running = True
        self.db_lock = threading.Lock()  # 创建一个锁来保护 self.db
        test_thread = threading.Thread(target=test_task)
        self.threads.append(test_thread)
        test_thread.start()

        # 导出接口文件

    def export_interface_file(self):
        # 定义输出顺序
        output_order = [
            "央视",
            "卫视",
            "地方",
            "港澳台",
            "卡通",
            "海外",
            "直播中国",
            "非实时",
        ]
        links = []
        for row in self.db.values():
            if row[6] != "" and float(row[6]) > 0:
                links.append(row)

        def custom_sort(row):
            try:
                col3_index = output_order.index(row[2])
            except ValueError:
                # 给一个较大的值，让不在sort_order中的元素排到最后
                col3_index = len(output_order)

            # 获取台名row[0]的第一个数字位置
            first_digit_index = next(
                (i for i, char in enumerate(row[0]) if char.isdigit()), None
            )
            # 从第一个数字开始截取字符串
            if first_digit_index is not None:
                name_start = row[0][:first_digit_index]
                name_end = row[0][first_digit_index:]
                if name_end and name_end.isdigit():
                    try:
                        name_end = int(name_end)
                    except ValueError:
                        name_end = 0
                else:
                    name_end = 0
            else:
                name_start = row[0]
                name_end = 0
            return col3_index, name_start, name_end, -float(row[6])

        result = sorted(links, key=custom_sort)
        with open("live.json", "w", encoding="utf-8") as file:
            file.write(
                "#权且苟同优选@www.ynlo.top:"
                + datetime.now().today().strftime("%Y-%m-%d %H:%M:%S" + "\n")
            )
            region = None
            for row in result:
                if row[2] != region:
                    region = row[2]
                    file.write(f"{region},#genre#\n")
                file.write(f"{row[0]},{row[1]}\n")
        self.console_print("导出接口文件成功!")
        # 打开导出文件
        os.startfile("live.json")
    def dev_test(self):
        for i in range(100):
            self.console_print(f"测试{i}")
            # 更新root
            time.sleep(0.1)
 


if __name__ == "__main__":
    root = tk.Tk()
    app = LiveTestApp(root)
    root.mainloop()
